package com.raven.play.datalink.evaluate;

import akka.Done;
import akka.stream.*;
import akka.stream.stage.AbstractInHandler;
import akka.stream.stage.AbstractOutHandler;
import akka.stream.stage.GraphStageLogic;
import akka.stream.stage.GraphStageWithMaterializedValue;
import org.codehaus.commons.compiler.CompilerFactoryFactory;
import org.codehaus.commons.compiler.IScriptEvaluator;
import scala.Tuple2;

import java.lang.reflect.ParameterizedType;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;

/**
 * 因为 I, O 是泛型, 如果不在子类中给出明确的类型, 是无法获得具体类型的, 而 Janino 模拟执行脚本必须知道入参, 出参类型
 * 所以将该类设置为 abstract class
 *
 * @param <I>
 * @param <O>
 */
public abstract class ScriptEvalFlowGraphStage<I, O> extends GraphStageWithMaterializedValue<FlowShape<I, O>, CompletionStage<Done>> {

    private final Inlet<I> inlet = Inlet.create("ScriptEvalFlow.in");
    private final Outlet<O> outlet = Outlet.create("ScriptEvalFlow.out");
    private final FlowShape<I, O> shape = new FlowShape<>(inlet, outlet);

    /**
     * 脚本运行使用的jar包
     */
    private final static String SUPPORT_LIB_DIR = "lib/";
    private final String script;

    private final Class<I> inletClass;
    private final Class<O> outletClass;

    @SuppressWarnings("unchecked")
    public ScriptEvalFlowGraphStage(String script) {
        this.script = script;

        this.inletClass = (Class<I>) ((ParameterizedType) getClass().getGenericSuperclass()).getActualTypeArguments()[0];
        this.outletClass = (Class<O>) ((ParameterizedType) getClass().getGenericSuperclass()).getActualTypeArguments()[1];
    }

    @Override
    public FlowShape<I, O> shape() {
        return shape;
    }

    @Override
    public Tuple2<GraphStageLogic, CompletionStage<Done>> createLogicAndMaterializedValue(Attributes inheritedAttributes) throws Exception {
        CompletableFuture<Done> promise = new CompletableFuture<>();
        return new Tuple2<>(new ScriptEvalFlowGraphStageLogic(shape, promise), promise);
    }

    class ScriptEvalFlowGraphStageLogic extends GraphStageLogic {

        private final CompletableFuture<Done> promise;
        private IScriptEvaluator scriptEvaluator;

        public ScriptEvalFlowGraphStageLogic(Shape shape, CompletableFuture<Done> promise) {
            super(shape);
            this.promise = promise;

            this.setHandler(inlet, new AbstractInHandler() {
                @Override
                public void onPush() {
                    try {
                        I input = grab(inlet);
                        O result = (O) scriptEvaluator.evaluate(new Object[]{input});
                        push(outlet, result);
                    } catch (Exception e) {
                        promise.completeExceptionally(e);
                        failStage(e);
                    }
                }

            });

            this.setHandler(outlet, new AbstractOutHandler() {
                @Override
                public void onPull() throws Exception {
                    pull(inlet);
                }
            });
        }

        @Override
        public void preStart() throws Exception {
            RestrictedClassLoader restrictedClassLoader = new RestrictedClassLoader(SUPPORT_LIB_DIR);
            scriptEvaluator = CompilerFactoryFactory
                    .getDefaultCompilerFactory(restrictedClassLoader)
                    .newScriptEvaluator();

            scriptEvaluator.setParameters(new String[]{"record"}, new Class[]{inletClass});
            scriptEvaluator.setReturnType(outletClass);

            scriptEvaluator.cook(script);
        }

        @Override
        public void postStop() throws Exception {
            if (!promise.isDone()) {
                promise.complete(Done.getInstance());
            }
        }

    }
}
